import asyncio, itertools, secrets, hashlib, typing as T
from playwright.async_api import async_playwright, Page

MERMAID_UMD = "https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.min.js"
# 如果你更喜欢 ESM，可改为 .esm.min.mjs + type="module" 注入（见下）

class MermaidRenderer:
    """一次启动浏览器，多次渲染共用，同步回输入顺序"""
    def __init__(self, workers: int = 4):
        self._workers = workers
        self._cnt = itertools.count()
        self._cache: set[str] = set()

    async def __aenter__(self):
        pw = await async_playwright().start()
        self._browser = await pw.chromium.launch()
        self._page: Page = await self._browser.new_page()

        # ① 注入 UMD 版本，直接得到 window.mermaid
        await self._page.add_script_tag(url=MERMAID_UMD)        # UMD 会挂 window.mermaid
        await self._page.wait_for_function("window.mermaid")    # 保守起见等待脚本 onload
        await self._page.evaluate("mermaid.initialize({startOnLoad:false});")
        return self

    async def render_one(self, code: str) -> dict:
        idx = next(self._cnt)
        res = {"parsed": False, "rendered": False, "error": None, "idx": idx}
        h = hashlib.sha1(code.encode()).hexdigest()
        if h in self._cache:            # 重复代码直接复用结果
            res.update(parsed=True, rendered=True)
            return res

        try:
            elem_id = f"m{idx}_{secrets.token_hex(4)}"
            svg: str = await self._page.evaluate(
                """({code, elem_id}) => mermaid.render(elem_id, code).then(r => r.svg)""",
                {"code": code, "elem_id": elem_id}
            )
            res.update(parsed=True, rendered=bool(svg))
            if svg:
                self._cache.add(h)
            else:
                res["error"] = "empty SVG"
        except Exception as e:
            res["error"] = str(e)
        return res

    async def __aexit__(self, *exc):
        await self._browser.close()

async def validate_codes(codes: list[str], workers: int = 4) -> list[dict]:
    async with MermaidRenderer(workers) as R:
        sem = asyncio.Semaphore(workers)
        async def task(code):
            async with sem:
                return await R.render_one(code)
        tasks = [task(c) for c in codes]
        results = await asyncio.gather(*tasks)
        # 结果还原为与输入一一对应的顺序
        ordered = [None] * len(codes)
        for r in results:
            ordered[r["idx"]] = r
            del r["idx"]
        ordered = [r.get("parsed", False) for r in ordered]
        return ordered



if __name__ == "__main__":
    src = [
        "graph TD; A-->B;",
        "sequenceDiagram\nA->>B: ok;",
        "graph LR; A-- B"              # 语法错
    ]

    datas = []
    mermaid_text = """graph TD
        start([开始]) --> preprocess[预处理对话文本]
        preprocess --> extractBasic[提取基本信息]
        extractBasic --> extractChief[提取主诉]
        extractChief --> extractHpi[提取现病史]
        extractHpi --> extractPmh[提取既往史]
        extractPmh --> extractPe[提取体格检查]
        extractPe --> fill[填充病历结构]
        fill --> end([结束])"""
    datas.append(mermaid_text)
    mermaid_text = """graph TD
        开始 --> 基本信息提取
        基本信息提取 --> 主诉识别
        主诉识别 --> 现病史分析
        现病史分析 --> 既往史判断
        既往史判断 --> 体格检查分析
        体格检查分析 --> 信息完整性验证
        信息完整性验证 -->|关键信息完整| 病历生成
        信息完整性验证 -->|关键信息缺失| 信息不足标记
        病历生成 --> 结束
        信息不足标记 --> 结束"""
    datas.append(mermaid_text)
    mermaid_text = """graph TD
        "开始" --> "基本信息提取"
        "基本信息提取" --> "主诉识别"
        "主诉识别" --> "现病史分析"
        "现病史分析" --> "既往史判断"
        "既往史判断" --> "体格检查分析"
        "体格检查分析" --> "信息完整性验证"
        "信息完整性验证" -->|"关键信息完整"| "病历生成"
        "信息完整性验证" -->|"关键信息缺失"| "信息不足标记"
        "病历生成" --> "结束"
        "信息不足标记" --> "结束" """
    datas.append(mermaid_text)

    # datas = datas * 4

    # 统计执行时间
    import time
    start_time = time.time()
    out = asyncio.run(validate_codes(datas, workers=3))
    end_time = time.time()
    print(f"验证耗时: {end_time - start_time:.4f} 秒")
    for row in out:
        print(row)
    print('End')
